// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once
#include <boost/algorithm/string.hpp>
#include "user_info.h"
#include "proto/common.pb.h"
#include "mem_row_descriptor.h"
#include "network_socket.h"
#include "table_record.h"
#include "runtime_state.h"
#include "base.h"
#include "expr.h"
// #include "range.h"
#include "fragment.h"

namespace baikaldb {
DECLARE_bool(default_2pc);

class ExecNode;

// notice日志信息统计结构
struct QueryStat {
    int64_t     query_read_time;
    int64_t     query_plan_time;
    int64_t     query_exec_time;
    int64_t     result_pack_time;
    int64_t     result_send_time;
    int64_t     unit_to_req_time;
    int64_t     req_to_buf_time;
    int64_t     server_talk_time;
    int64_t     buf_to_res_time;
    int64_t     res_to_table_time;
    int64_t     table_get_row_time;
    int64_t     total_time;
    uint64_t    version;
    size_t      send_buf_size;
    int32_t     partition_key;
    
    int32_t     sql_length;
    int32_t     region_count;
    bool        hit_cache;
    bool        hit_query_cache;
    timeval     start_stamp;
    timeval     send_stamp;
    timeval     end_stamp;
    //std::string traceid;
    std::string family;
    std::string table;
    std::string resource_tag;
    std::string server_ip;
    std::ostringstream sample_sql;
    std::string trace_id;
    uint64_t    sign = 0;
    int64_t     table_id = -1;

    MysqlErrCode error_code;
    std::ostringstream error_msg;

    int64_t     num_affected_rows = 0;
    int64_t     num_returned_rows = 0;
    int64_t     num_scan_rows     = 0;
    int64_t     read_disk_size = 0;
    int64_t     num_filter_rows   = 0;
    int64_t     txn_alive_time    = 0;
    uint64_t    log_id = 0;
    uint64_t    old_txn_id = 0;
    int         old_seq_id = 0;

    int64_t     db_handle_rows    = 0;
    int64_t     db_handle_bytes   = 0;

    QueryStat() {
        reset();
    }

    void reset() {
        query_read_time     = 0;
        query_plan_time     = 0;
        query_exec_time     = 0;
        result_pack_time    = 0;
        result_send_time    = 0;
        unit_to_req_time    = 0;
        req_to_buf_time     = 0;
        server_talk_time    = 0;
        buf_to_res_time     = 0;
        res_to_table_time   = 0;
        table_get_row_time  = 0;
        total_time          = 0;
        version             = 0;
        send_buf_size       = 0;
        partition_key       = 0;
        sql_length          = 0;
        txn_alive_time      = 0;
        hit_cache           = false;
        hit_query_cache     = false;
        gettimeofday(&(start_stamp), NULL);
        gettimeofday(&(send_stamp), NULL);
        gettimeofday(&(end_stamp), NULL);
        //traceid.clear();
        family.clear();
        table.clear();
        table_id = -1;
        server_ip.clear();
        sample_sql.str("");
        trace_id.clear();
        sign = 0;

        error_code          = ER_ERROR_FIRST;
        error_msg.str("");
        num_affected_rows   = 0;
        num_returned_rows   = 0;
        num_scan_rows       = 0;
        read_disk_size      = 0;
        num_filter_rows     = 0;
        log_id              = butil::fast_rand();
        old_txn_id          = 0;
        old_seq_id          = 0;
        region_count        = 0;
        db_handle_rows      = 0;
        db_handle_bytes     = 0;
    }
};

struct ExprParams {
    bool is_expr_subquery = false;
    bool is_correlated_subquery = false;
    bool is_from_subquery = false;
    parser::FuncType    func_type = parser::FT_COMMON;
    parser::CompareType cmp_type  = parser::CMP_ANY;
    // (a,b) in (select a,b from t) row_filed_number=2
    int  row_filed_number  = 1;
    bool is_union_subquery = false;
};

struct UniqueIdContext {
    int32_t tuple_cnt = 0;
    int32_t sub_query_level = 0;
    int32_t derived_table_id = -1;
};
typedef std::shared_ptr<UniqueIdContext> SmartUniqueIdCtx;

// explain format='hint WITH_CBO IGNORE_INDEX(index_name1,index_name2)'
// 用以explain时给提示
// NO_FORCE 会在索引选择时忽视force标记
// NO_USE 会在索引选择是忽视use标记
// WITH_CBO 会在索引选择时使用cbo选择索引
// WITHOUT_CBO 会在索引选择是不使用cbo选择索引
// IGNORE_INDEX 可以在索引选择时屏蔽某些索引 IGNORE_INDEX(index_name1,index_name2) indexname之间不要有空格
// NO_IGNORE 会在索引选择是忽视ignore标记
// IGNORE_BLACKLIST 允许被加黑的sql 进行explain
struct ExplainHint {
    enum HintType {
        NO_FORCE,
        NO_USE,
        NO_IGNORE,
        WITH_CBO,
        WITHOUT_CBO,
        IGNORE_INDEX,
        KEEP_VIRTUAL,
    };

    explicit ExplainHint(std::string format) {
        std::vector<std::string> hints;
        boost::split(hints, format, boost::is_any_of(" \t\n\r"), boost::token_compress_on);
        for (const auto& hint: hints) {
            if (boost::iequals(hint, "NO_FORCE")) {
                set_flag<NO_FORCE>();
            } else if (boost::iequals(hint, "NO_USE")) {
                set_flag<NO_USE>();
            } else if (boost::iequals(hint, "WITH_CBO")) {
                set_flag<WITH_CBO>();
            } else if (boost::iequals(hint, "WITHOUT_CBO")) {
                set_flag<WITHOUT_CBO>();
            } else if (boost::iequals(hint, "NO_IGNORE")) {
                set_flag<NO_IGNORE>();
            } else if (boost::iequals(hint, "KEEP_VIRTUAL")) {
                set_flag<KEEP_VIRTUAL>();
            } else if (boost::istarts_with(hint, "IGNORE_INDEX")) {
                set_flag<IGNORE_INDEX>();
                // 提取index names
                std::string IGNORE_INDEX_names = hint.substr(12, hint.length() - 13);
                std::vector<std::string> IGNORE_INDEX_names_list;
                boost::split(IGNORE_INDEX_names_list, IGNORE_INDEX_names, boost::is_any_of(","), boost::token_compress_on);
                blocked_indexes.insert(blocked_indexes.end(), IGNORE_INDEX_names_list.begin(), IGNORE_INDEX_names_list.end());
            }
        }
    }

    template<HintType ht>
    void set_flag() {
        constexpr long long mask = (1ULL << ht);
        hint_bit_map |= mask;
    }

    template<HintType ht>
    void unset_flag() {
        constexpr long long mask = ~(1ULL << ht);
        hint_bit_map &= mask;
    }

    template<HintType ht>
    bool get_flag() {
        return hint_bit_map & (1ULL << ht);
    }

    long long hint_bit_map = 0;
    std::vector<std::string> blocked_indexes;
};

class QueryContext {
public:
    QueryContext() {
        enable_2pc = FLAGS_default_2pc;
    }
    QueryContext(std::shared_ptr<UserInfo> user, std::string db) : 
            cur_db(db),
            user_info(user) {
        enable_2pc = FLAGS_default_2pc;
    }

    ~QueryContext();

    void add_tuple(const pb::TupleDescriptor& tuple_desc) {
        if (tuple_desc.tuple_id() >= (int)_tuple_descs.size()) {
            _tuple_descs.resize(tuple_desc.tuple_id() + 1);
        }
        _tuple_descs[tuple_desc.tuple_id()] = tuple_desc;
    }

    pb::TupleDescriptor* get_tuple_desc(int tuple_id) {
        if (tuple_id < 0 || tuple_id >= (int)_tuple_descs.size()) {
            return nullptr;
        }
        return &_tuple_descs[tuple_id];
    }
    std::vector<pb::TupleDescriptor>* mutable_tuple_descs() {
        return &_tuple_descs;
    }
    const std::vector<pb::TupleDescriptor>& tuple_descs() {
        return _tuple_descs;
    }
/*
    int32_t get_tuple_id(int64_t table_id) {
        for (auto& tuple_desc : _tuple_descs) {
            if (tuple_desc.table_id() == table_id) {
                return tuple_desc.tuple_id();
            }
        }
        return -1;
    }
*/
    int32_t get_slot_id(int32_t tuple_id, int32_t field_id) {
        for (const auto& slot_desc : _tuple_descs[tuple_id].slots()) {
            if (slot_desc.field_id() == field_id) {
                return slot_desc.slot_id();
            }
        }
        return -1;
    }
    SmartState get_runtime_state() {
        if (runtime_state == nullptr) {
            runtime_state.reset(new RuntimeState);
        }
        return runtime_state;
    }
    pb::PlanNode* add_plan_node() {
        return plan.add_nodes();
    }
    int create_plan_tree();
    int destroy_plan_tree();

    void add_sub_ctx(std::shared_ptr<QueryContext>& ctx) {
        std::unique_lock<bthread::Mutex> lck(_kill_lock);
        sub_query_plans.emplace_back(ctx);
    }

    void add_noncorrelated_subquerys(std::shared_ptr<QueryContext>& ctx) {
        std::unique_lock<bthread::Mutex> lck(_kill_lock);
        noncorrelated_subquerys.emplace_back(ctx);
    }

    void set_kill_ctx(std::shared_ptr<QueryContext>& ctx) {
        std::unique_lock<bthread::Mutex> lck(_kill_lock);
        kill_ctx = ctx;
    }
    void kill_all_ctx() {
        std::unique_lock<bthread::Mutex> lck(_kill_lock);
        get_runtime_state()->cancel();
        for (auto& ctx : sub_query_plans) {
            ctx->get_runtime_state()->cancel();
        }
        if (kill_ctx) {
            kill_ctx->get_runtime_state()->cancel();
        }
    }
    void update_ctx_stat_info(RuntimeState* state, int64_t total_time);
    
    int64_t get_ctx_total_time();

    bool need_use_read_index() {
        if (user_info == nullptr) {
            return false;
        }
        return user_info->need_use_read_index();
    }

    // 用于PreparePlanner和PlanCache复制QueryContext
    int copy_query_context(QueryContext* p_query_context);

    // 用于最外层查询获取所有的子查询QueryContext
    void get_all_derived_table_ctx_mapping(std::map<int32_t, std::shared_ptr<QueryContext>>& m) {
        for (auto& kv : derived_table_ctx_mapping) {
            m[kv.first] = kv.second;
            if (kv.second != nullptr) {
                kv.second->get_all_derived_table_ctx_mapping(m);
            }
        }
    }
    // 用于最外层查询获取所有的子查询slot_column_mapping
    void get_all_slot_column_mapping(std::map<int32_t, std::map<int32_t, int32_t>>& m) {
        for (auto& kv : slot_column_mapping) {
            m[kv.first] = kv.second;
        }
        for (auto& kv : derived_table_ctx_mapping) {
            if (kv.second != nullptr) {
                kv.second->get_all_slot_column_mapping(m);
            }
        }
    }

    // only use mpp
    void set_client_conn(NetworkSocket* conn) {
        client_conn = conn;
    }
public:
    std::string         sql;
    std::vector<std::string> comments;
    std::string         cur_db;
    pb::Charset         charset = pb::CS_UNKNOWN; // Connection charset
    pb::TraceNode       trace_node;

    // new sql parser data structs
    parser::StmtNode*   stmt;
    parser::NodeType    stmt_type;
    bool                is_explain = false;
    std::shared_ptr<ExplainHint> explain_hint;
    std::vector<std::shared_ptr<QueryContext>> noncorrelated_subquerys;
    bool                is_get_keypoint = false;
    bool                is_full_export = false;
    bool                is_straight_join = false;
    bool                select_for_update = false;
    ExplainType         explain_type = EXPLAIN_NULL;
    // Explain format内的具体信息
    std::string         format;
    int                 single_store_concurrency = -1;

    uint8_t             mysql_cmd = COM_SLEEP;      // Command number in mysql protocal.
    int                 type;           // Query type. finer than mysql_cmd.
    int64_t             row_ttl_duration = 0; // used for /*{"duration": xxx}*/ insert ...
    uint64_t            watt_stats_version = 0; // used for /*{"watt_stats_version": xxx}*/ insert merge into ...
    QueryStat           stat_info;      // query execute result status info
    std::shared_ptr<UserInfo> user_info;

    pb::Plan            plan;
    ExecNode*           root = nullptr;

    // mpp TODO CACHE?
    std::shared_ptr<FragmentInfo>        root_fragment;   // fragment树
    std::map<int, std::shared_ptr<FragmentInfo>> fragments;
    std::map<std::string, std::set<int>> db_to_fragments; // 非主db address0->fragment_ids
    int                 fragment_max_id = 0;
    int                 node_max_id = 0;                  // 只有Exchange有nodeid,一对Exchange的nodeid相同
    
    // 如果输出包含AGG表达式，则在PacketNode和AggNode节点中的常量节点都需要获取并赋值
    std::unordered_multimap<int, ExprNode*> placeholders;
    std::string         prepare_stmt_name;
    std::vector<pb::ExprNode> param_values;

    SmartState          runtime_state;  // baikaldb side runtime state
    NetworkSocket*      client_conn = nullptr; // used for baikaldb
    // the insertion records, not grouped by region yet
    std::vector<SmartRecord>            insert_records;

    bool                succ_after_logical_plan = false;
    bool                succ_after_physical_plan = false;
    bool                return_empty = false;
    bool                new_prepared = false;  // flag for stmt_prepare
    bool                exec_prepared = false; // flag for stmt_execute
    bool                is_prepared = false;   // flag for stmt_execute
    bool                is_select = false;
    bool                need_destroy_tree = false;
    bool                has_derived_table = false;
    bool                has_information_schema = false;
    bool                is_complex = false;
    bool                use_backup = false;
    bool                need_learner_backup = false;
    int64_t             prepared_table_id = -1;
    ExprParams          expr_params;
    // field: column_id
    std::map<std::string, int32_t> field_column_id_mapping;
    // tuple_id: field: slot_id
    std::map<int64_t, std::map<std::string, int32_t>> ref_slot_id_mapping;
    // tuple_id: slot_id: column_id
    std::map<int64_t, std::map<int32_t, int32_t>> slot_column_mapping;
    std::map<int64_t, std::shared_ptr<QueryContext>> derived_table_ctx_mapping;
    // 当前sql涉及的所有tuple
    std::set<int64_t>   current_tuple_ids;
    // 当前sql涉及的表的tuple
    std::set<int64_t>   current_table_tuple_ids;
    bool                open_binlog = false;
    bool                no_binlog = false; // 用于控制DM导入是否写binlog
    SignExecType        sql_exec_type_defined = SignExecType::SIGN_EXEC_NOT_SET;
    bool                use_mpp = false;

    // user can scan data in specific region by comments 
    // /*{"region_id":$region_id}*/ preceding a Select statement 
    int64_t             debug_region_id = -1;

    // user can scan data in specific peer by comments
    // /*{"peer_index":$peer_index}*/ preceding a Select statement
    int64_t             peer_index = -1;

    // user can get query data from cache in specific query_cache(ms) by comments
    // /*{"query_cache":$query_cache}*/ preceding a Select statement
    int64_t             query_cache = 0;

    // in autocommit mode, two phase commit is disabled by default (for better formance)
    // user can enable 2pc by comments /*{"enable_2pc":1}*/ preceding a DML statement
    bool                enable_2pc = false;
    bool                is_cancelled = false;
    bool                execute_global_flow = false;
    std::shared_ptr<QueryContext> kill_ctx;
    std::vector<std::shared_ptr<QueryContext>> sub_query_plans;
    std::unordered_map<uint64_t, std::string> long_data_vars;
    std::vector<SignedType> param_type;
    std::set<int64_t> index_ids;
    // 用于索引推荐
    std::map<int32_t, int> field_range_type;
    std::set<uint64_t> sign_blacklist;
    std::set<uint64_t> sign_forcelearner;
    std::set<uint64_t> sign_rolling;
    std::map<uint64_t, std::map<int64_t, std::set<std::string>>> sign_forceindex; // sign => <table_id, index_name_set>

    std::map<int64_t, std::vector<std::string>> table_partition_names;
    std::map<uint64_t, SignExecType> sign_exec_type;

    std::map<std::string, std::string> table_with_clause_mapping;

    // for base subscribe
    bool                is_base_subscribe = false;
    int64_t             base_subscribe_table_id = -1;
    std::string         base_subscribe_table_name = "";
    std::string         base_subscribe_filter_field = "";
    std::unordered_map<std::string, std::string> base_subscribe_select_name_alias_map {};

    // non-prepare plan cache
    bool is_plan_cache = false;
    bool has_find_placeholder = false;
    PlanCacheKey cache_key;
    std::unordered_map<int64_t, int64_t> table_version_map;
    bool has_unable_cache_expr = false;

    // 单表编码转换
    bool need_convert_charset = false;
    pb::Charset table_charset = pb::CS_UNKNOWN;

    // 是否为from型子查询
    bool is_from_subquery = false;
    int  sub_query_level = 0;

    // 是否为Union型子查询
    bool is_union_subquery = false;

    // 是否为insert select里的select子查询
    bool is_insert_select_subquery = false;

    bool table_can_use_arrow_vectorize = true;
    // 是否是with语句
    bool is_with = false;
    // 是否是create view语句
    bool is_create_view = false;

    // 是否包含窗口函数
    bool has_window_func = false;

    // 聚合下推使用
    SmartUniqueIdCtx unique_id_ctx;

    // 向量索引使用
    int32_t efsearch = -1;

    bool dumped_slow_sql = false;

    // 注释或者配置指定的sign mpp hash值
    int mpp_hash_num = -1;

    // 查询是否包含DBLink Mysql表
    bool has_dblink_mysql = false;

private:
    std::vector<pb::TupleDescriptor> _tuple_descs;
    bthread::Mutex _kill_lock;
};

inline int32_t get_slot_idx(const pb::TupleDescriptor& tuple_desc, const int32_t slot_id) {
    if (tuple_desc.slot_idxes().empty()) {
        // 兼容MPP场景未上线的db还未赋值slot_idxes，将任务发送到已上线的db
        return slot_id - 1;
    }
    const auto& slot_idxes = tuple_desc.slot_idxes();
    if (slot_id - 1 < 0 || slot_id - 1 >= slot_idxes.size()) {
        DB_WARNING("Invalid slot_id: %d, slot_idxes size: %d", slot_id, slot_idxes.size());
        return -1;
    }
    return slot_idxes[slot_id - 1];
}

} //namespace baikal
